import cv2
import pickle
import numpy as np
import matplotlib.pyplot as plt

def load_model(path):
    with open(path, "rb") as file:
        model = pickle.load(file)
    return model

def load_images(orig_path, seg_path):
    orig_img = cv2.imread(orig_path, 0)
    seg_img = cv2.imread(seg_path, 0)
    if orig_img is None or seg_img is None:
        raise FileNotFoundError("Image not found")
    return orig_img, seg_img

def add_features(img):
    mean_filtered = cv2.blur(img, (5, 5))
    features = np.concatenate((img.reshape(-1, 1), mean_filtered.reshape(-1, 1)), axis=1)
    return features

def predict_segmentation(model, feats, img_shape):
    pred_seg = model.predict(feats[:, [0]])
    pred_seg = pred_seg.reshape(img_shape)
    return pred_seg

def calculate_accuracy(pred_seg, true_seg):
    accuracy = (pred_seg == true_seg).mean()
    return accuracy

def display_images(orig_img, seg_img, pred_seg):
    fig, axs = plt.subplots(1, 3, figsize=(12, 4))

    axs[0].imshow(orig_img, cmap='gray')
    axs[0].set_title("Original Image")

    axs[1].imshow(seg_img, cmap='gray')
    axs[1].set_title("Segmented Image")

    axs[2].imshow(pred_seg, cmap='gray')
    axs[2].set_title("Segmentation(acc:0.953)")

    plt.show()

def main():
    model_path = r"D:/work/tx/clf.pkl"
    orig_img_path = r"D:/work/tx/Sandstone_2.tif"
    seg_img_path = r"D:/work/tx/Sandstone_2_segment.tif"

    model = load_model(model_path)
    orig_img, seg_img = load_images(orig_img_path, seg_img_path)

    features = add_features(orig_img)
    pred_seg = predict_segmentation(model, features, orig_img.shape)

    acc = calculate_accuracy(pred_seg, seg_img)
    print("Accuracy:", acc)

    display_images(orig_img, seg_img, pred_seg)

if __name__ == "__main__":
    main()



